
package sf.r2dbc.dao;

import io.r2dbc.spi.Connection;
import io.r2dbc.spi.ConnectionFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import sf.common.wrapper.Page;
import sf.core.DBCascadeField;
import sf.core.DBObject;
import sf.database.dao.DBContext;
import sf.database.dbinfo.DBMetaData;
import sf.database.dialect.DBDialect;
import sf.database.jdbc.sql.PageStrategy;
import sf.database.jdbc.sql.SQLContext;
import sf.database.jdbc.sql.SQLParameter;
import sf.database.meta.MetaHolder;
import sf.database.meta.TableMapping;
import sf.dsl.DeleteDSL;
import sf.dsl.InsertDSL;
import sf.dsl.UpdateDSL;
import sf.dsl.example.Example;
import sf.r2dbc.binding.BindMarkersFactory;
import sf.r2dbc.connection.R2dbcConnection;
import sf.r2dbc.sql.R2dbcContext;
import sf.r2dbc.sql.R2dbcCrud;
import sf.r2dbc.sql.R2dbcUtils;
import sf.spring.util.Assert;

import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * r2dbc的sorm适配
 * 该适配兼容jdbc模式,可以使用?占位符,执行时会自动替换?占位符.
 * @see org.springframework.data.r2dbc.core.R2dbcEntityTemplate
 */
public class R2dbcClient implements R2dbcSupport {
    protected ConnectionFactory connectionFactory;
    protected DBDialect dbDialect;
    protected R2dbcCrud crud = R2dbcCrud.getInstance();

    /**
     * sql执行上下文.使用Mono/Flux中的ContextView传递上下文时,存在全部生效的问题,无法精细控制生效时机.
     */
    protected ThreadLocal<DBContext> dbContexts = ThreadLocal.withInitial(() -> null);

    public R2dbcClient(ConnectionFactory connectionFactory) {
        Assert.notNull(connectionFactory, "ConnectionFactory must not be null");
        this.connectionFactory = connectionFactory;
        dbDialect = R2dbcUtils.getDBDialect(connectionFactory);
    }

    @Override
    public R2dbcClient useContext(DBContext context) {
        Assert.notNull(context, "Context is null!");
        if (context.isUseLazyLoad()) {
            context.setClient(this);
        }
        dbContexts.set(context);
        return this;
    }

    protected DBContext getDBContext() {
        DBContext context = dbContexts.get();
        if (context != null) {
            dbContexts.remove();
        }
        return context;
    }

    protected R2dbcConnection getR2dbcConnection(Connection connection, DBContext dbContext) {
        R2dbcConnection w = new R2dbcConnection(connection, connectionFactory);
        BindMarkersFactory bmf = R2dbcUtils.getBindMarkersFactory(connection);
        R2dbcContext r2dbcContext = new R2dbcContext();
        r2dbcContext.setBindMarkersFactory(bmf);
        w.setR2dbcContext(r2dbcContext);
        w.setDbDialect(dbDialect);
        w.setDbContext(dbContext);
        return w;
    }

    @Override
    public Mono<DBDialect> getDialect() {
        return monoFunc(r2dbcConnection -> Mono.just(R2dbcUtils.getDBDialect(r2dbcConnection)));
    }

    @Override
    public <T> Flux<T> fluxFunc(Function<Connection, Flux<T>> action) {
        DBContext dbContext = getDBContext();
        return Mono.from(connectionFactory.create()).flatMapMany(connection -> {
            R2dbcConnection r2dbcConnection = getR2dbcConnection(connection, dbContext);
            return action.apply(r2dbcConnection);
        });
    }

    @Override
    public <T> Mono<T> monoFunc(Function<Connection, Mono<T>> action) {
        DBContext dbContext = getDBContext();
        return Mono.from(connectionFactory.create()).flatMap(connection -> {
            R2dbcConnection r2dbcConnection = getR2dbcConnection(connection, dbContext);
            return action.apply(r2dbcConnection);
        });
    }

    @Override
    public <T extends DBObject> Mono<Long> selectCount(T query) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().selectCount(r2dbcConnection, query));
    }

    @Override
    public <T extends DBObject> Mono<Boolean> existsByPrimaryKeys(Class<T> clz, Object... keyParams) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().existsByPrimaryKeys(r2dbcConnection, clz, keyParams));
    }

    @Override
    public <T extends DBObject> Mono<Boolean> exists(T query) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().exists(r2dbcConnection, query));
    }

    @Override
    public <T extends DBObject> Flux<T> selectList(T query) {
        return fluxFunc(r2dbcConnection -> crud.getCrudModel().selectList(r2dbcConnection, query));
    }

    @Override
    public <T extends DBObject> Flux<T> selectList(T query, long start, int limit) {
        return fluxFunc(r2dbcConnection -> crud.getCrudModel().selectList(r2dbcConnection, query));
    }

    @Override
    public <T extends DBObject> Flux<T> selectListForUpdate(T query) {
        return fluxFunc(r2dbcConnection -> crud.getCrudModel().selectListForUpdate(r2dbcConnection, query));
    }

    @Override
    public <T extends DBObject> Mono<Page<T>> selectPage(T query, long start, int limit) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().selectPage(r2dbcConnection, query, start, limit));
    }

    @Override
    public <T> Mono<Page<T>> selectPage(Class<T> clz, SQLContext context, long start, int limit) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().selectPage(r2dbcConnection, start, limit, clz, context));
    }

    @Override
    public <T extends DBObject> Mono<T> selectByPrimaryKeys(Class<T> clz, Object... keyParams) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().selectByPrimaryKeys(r2dbcConnection, clz, keyParams));
    }

    @Override
    public <T extends DBObject> Mono<T> selectOne(T query) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().selectOne(r2dbcConnection, query));
    }

    @Override
    public <T extends DBObject> Mono<Integer> insert(T obj) {
        return monoFunc(r2dbcConnection ->
                crud.getCrudModel().insert(r2dbcConnection, obj, false, true));
    }

    @Override
    public <T extends DBObject> Mono<Integer> insert(T obj, boolean fast, boolean useOptimisticLock) {
        return monoFunc(r2dbcConnection ->
                crud.getCrudModel().insert(r2dbcConnection, obj, fast, useOptimisticLock));
    }

    @Override
    public <T extends DBObject> Mono<Integer> update(T entity) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().update(r2dbcConnection, entity));
    }

    @Override
    public <T extends DBObject> Mono<Integer> updateAndSet(T obj) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().updateAndSet(r2dbcConnection, obj));
    }

    @Override
    public Mono<Integer> merge(DBObject obj) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().merge(r2dbcConnection, obj));
    }

    @Override
    public <T extends DBObject> Mono<Integer> deleteByPrimaryKeys(Class<T> clz, Object... keyParams) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().deleteByPrimaryKeys(r2dbcConnection, clz, keyParams));
    }

    @Override
    public <T extends DBObject> Mono<Integer> delete(T entity) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().delete(r2dbcConnection, entity));
    }

    @Override
    public <T extends DBObject> Mono<Integer> logicDeleteByPrimaryKeys(Class<T> clz, Object... keyParams) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().logicDeleteByPrimaryKeys(r2dbcConnection, clz, keyParams));
    }

    @Override
    public Mono<Integer> logicDelete(DBObject obj) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().logicDelete(r2dbcConnection, obj));
    }

    @Override
    public <T extends DBObject> Mono<Integer> truncate(Class<T> clz) {
        return monoFunc(r2dbcConnection -> truncate(r2dbcConnection, MetaHolder.getMeta(clz)));
    }

    /**
     * 执行 truncate table XXXX 命令，快速清光表的数据
     * @param conn
     * @param meta 表名
     * @return
     */
    public Mono<Integer> truncate(Connection conn, TableMapping meta) {
        DBDialect dialect = R2dbcUtils.getDBDialect(conn);
        List<String> sqls = DBMetaData.truncateSql(dialect, Collections.singletonList(meta.getTableName()));
        Flux<Integer> flux = null;
        for (String sql : sqls) {
            if (flux == null) {
                flux = crud.getCrudSql().execute(conn, sql).flux();
            } else {
                flux = flux.mergeWith(crud.getCrudSql().execute(conn, sql));
            }
        }
        if (flux != null) {
            return flux.collect(Collectors.summingInt(Number::intValue));
        }
        return Mono.just(0);
    }

    @Override
    public Flux<Integer> batchInsert(Collection<? extends DBObject> modelList) {
        return fluxFunc(r2dbcConnection ->
                crud.getCrudModel().batchInsert(r2dbcConnection, modelList, false, true));
    }

    @Override
    public Flux<Integer> batchInsertFast(Collection<? extends DBObject> modelList) {
        return fluxFunc(r2dbcConnection ->
                crud.getCrudModel().batchInsert(r2dbcConnection, modelList, true, true));
    }

    @Override
    public Flux<Integer> batchUpdate(Collection<? extends DBObject> modelList) {
        return fluxFunc(r2dbcConnection ->
                crud.getCrudModel().batchUpdate(r2dbcConnection, modelList));
    }

    @Override
    public Flux<Integer> batchDelete(Collection<? extends DBObject> modelList) {
        return fluxFunc(r2dbcConnection -> crud.getCrudModel().batchDelete(r2dbcConnection, modelList));
    }

    @Override
    public <T extends DBObject> Flux<T> fetchCascade(T query, Class<T> clz, DBCascadeField... fields) {
        return fluxFunc(r2dbcConnection -> crud.getCrudModel().fetchCascade(r2dbcConnection, query, clz, fields));
    }

    @Override
    public <T extends DBObject> Mono<T> fetchLinks(T obj) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().fetchLinks(r2dbcConnection, obj));
    }

    @Override
    public <T extends DBObject> Mono<T> fetchLinks(T obj, DBCascadeField... fields) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().fetchLinks(r2dbcConnection, obj, fields));
    }

    @Override
    public Mono<Integer> execute(SQLContext context) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().execute(r2dbcConnection, context));
    }

    @Override
    public Mono<Integer> insertCascade(DBObject obj, DBCascadeField... fields) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().insertCascade(r2dbcConnection, obj, fields));
    }

    @Override
    public <T extends DBObject> Mono<T> insertLinks(T obj, DBCascadeField... fields) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().insertLinks(r2dbcConnection, obj, fields));
    }

    @Override
    public <T extends DBObject> Mono<T> insertRelation(T obj, DBCascadeField... fields) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().insertRelation(r2dbcConnection, obj, fields));
    }

    @Override
    public Mono<Integer> updateCascade(DBObject obj, DBCascadeField... fields) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().updateCascade(r2dbcConnection, obj, fields));
    }

    @Override
    public <T extends DBObject> Mono<Integer> updateLinks(T obj, DBCascadeField... fields) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().updateLinks(r2dbcConnection, obj, fields));
    }

    @Override
    public <T extends DBObject> Mono<Integer> updateRelation(T obj, DBCascadeField... fields) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().updateRelation(r2dbcConnection, obj, fields));
    }

    @Override
    public <T extends DBObject> Mono<Integer> deleteCascade(T obj, DBCascadeField... fields) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().deleteCascade(r2dbcConnection, obj, fields));
    }

    @Override
    public <T extends DBObject> Mono<Integer> deleteLinks(T obj, DBCascadeField... fields) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().deleteLinks(r2dbcConnection, obj, fields));
    }

    @Override
    public <T extends DBObject> Mono<Integer> deleteRelation(T obj, DBCascadeField... fields) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().deleteRelation(r2dbcConnection, obj, fields));
    }

    @Override
    public <T> Mono<Page<T>> selectPage(long start, int limit, Class<T> beanClass, Example example) {
        return monoFunc(r2dbcConnection -> crud.getCrudExampleInf().selectPage(r2dbcConnection, start, limit, beanClass, example));
    }

    @Override
    public Mono<Object[]> selectArray(Example example) {
        return monoFunc(r2dbcConnection -> crud.getCrudExampleInf().selectArray(r2dbcConnection, example));
    }

    @Override
    public <T> Flux<T> selectList(Class<T> beanClass, Example example) {
        return fluxFunc(r2dbcConnection -> crud.getCrudExampleInf().selectList(r2dbcConnection, beanClass, example));
    }

    @Override
    public <T> Flux<T> selectList(Class<T> beanClass, long start, int limit, Example example) {
        return fluxFunc(r2dbcConnection -> crud.getCrudExampleInf().selectList(r2dbcConnection, beanClass, start, limit, example));
    }

    @Override
    public <T> Flux<T> selectListForUpdate(Class<T> beanClass, Example example) {
        return fluxFunc(r2dbcConnection -> crud.getCrudExampleInf().selectListForUpdate(r2dbcConnection, beanClass, example));
    }

    @Override
    public <T> Mono<T> selectOne(Class<T> beanClass, Example example) {
        return monoFunc(r2dbcConnection -> crud.getCrudExampleInf().selectOne(r2dbcConnection, beanClass, example));
    }

    @Override
    public <T> Mono<T> selectOneForUpdate(Class<T> beanClass, Example example) {
        return monoFunc(r2dbcConnection -> crud.getCrudExampleInf().selectOneForUpdate(r2dbcConnection, beanClass, example));
    }

    @Override
    public Flux<Map<String, Object>> select(Example example) {
        return fluxFunc(r2dbcConnection -> crud.getCrudExampleInf().select(r2dbcConnection, example));
    }

    @Override
    public Mono<Integer> insert(InsertDSL insertDSL, List<DBObject> datas) {
        return monoFunc(r2dbcConnection -> crud.getCrudExampleInf().insert(r2dbcConnection, insertDSL, datas));
    }

    @Override
    public Mono<Integer> update(UpdateDSL updateDSL) {
        return monoFunc(r2dbcConnection -> crud.getCrudExampleInf().update(r2dbcConnection, updateDSL));
    }

    @Override
    public Mono<Integer> delete(DeleteDSL deleteDSL) {
        return monoFunc(r2dbcConnection -> crud.getCrudExampleInf().delete(r2dbcConnection, deleteDSL));
    }

    @Override
    public Flux<Map<String, Object>> selectList(String sql, List<SQLParameter> paras) {
        return fluxFunc(r2dbcConnection -> crud.getCrudSql().select(r2dbcConnection, sql, paras));
    }

    @Override
    public <T> Mono<T> selectOne(Class<T> beanClass, String sql, List<SQLParameter> parameters) {
        return monoFunc(r2dbcConnection -> crud.getCrudSql().selectOne(r2dbcConnection, beanClass, sql, parameters));
    }

    @Override
    public <T extends DBObject> Mono<T> selectOneForUpdate(T query) {
        return monoFunc(r2dbcConnection -> crud.getCrudModel().selectOneForUpdate(r2dbcConnection, query));
    }

    @Override
    public <T> Flux<T> selectList(Class<T> beanClass, long start, int limit, String sql, List<SQLParameter> parameters) {
        return fluxFunc(r2dbcConnection -> crud.getCrudSql().selectList(r2dbcConnection, beanClass, start, limit, sql, parameters));
    }

    @Override
    public <T> Flux<T> selectList(Class<T> beanClass, String sql, List<SQLParameter> parameters) {
        return fluxFunc(r2dbcConnection -> crud.getCrudSql().selectList(r2dbcConnection, beanClass, sql, parameters));
    }

    @Override
    public <T> Mono<Page<T>> selectPage(long start, int limit, Class<T> beanClass, String sql, Object... parameters) {
        return monoFunc(r2dbcConnection -> crud.getCrudSql().selectPage(r2dbcConnection, start, limit, beanClass, sql, parameters));
    }

    @Override
    public <T> Mono<Page<T>> selectPageRaw(long start, int limit, Class<T> beanClass, String countSql, Object[] countParas, String listSql, Object[] listParas) {
        return monoFunc(r2dbcConnection -> crud.getCrudSql().selectPageRaw(r2dbcConnection, start, limit, beanClass, countSql, countParas, listSql, listParas));
    }

    @Override
    public <T> Mono<Page<T>> selectPageRaw(long start, int limit, Class<T> beanClass, String countSql, Object[] countParas, String listSql, Object[] listParas, PageStrategy strategy) {
        return monoFunc(r2dbcConnection -> crud.getCrudSql().selectPageRaw(r2dbcConnection, start, limit, beanClass, countSql, countParas, listSql, listParas, strategy));
    }

    @Override
    public Mono<Integer> insert(List<String> pkeys, Map<String, Object> keyValues, String sql, List<SQLParameter> parameters) {
        return monoFunc(r2dbcConnection -> crud.getCrudSql().insert(r2dbcConnection, pkeys, keyValues, sql, parameters));
    }

    @Override
    public Mono<Integer> execute(String sql, List<SQLParameter> paras) {
        return monoFunc(r2dbcConnection -> crud.getCrudSql().execute(r2dbcConnection, sql, paras));
    }

    @Override
    public Flux<Map<String, Object>> selectList(String sql, Object... parameters) {
        return fluxFunc(r2dbcConnection -> crud.getCrudSql().select(r2dbcConnection, sql, parameters));
    }

    @Override
    public <T> Mono<T> selectOne(Class<T> beanClass, String sql, Object... parameters) {
        return monoFunc(r2dbcConnection -> crud.getCrudSql().selectOne(r2dbcConnection, beanClass, sql, parameters));
    }

    @Override
    public <T> Flux<T> selectList(Class<T> beanClass, long start, int limit, String sql, Object... parameters) {
        return fluxFunc(r2dbcConnection -> crud.getCrudSql().selectList(r2dbcConnection, beanClass, start, limit, sql, parameters));
    }

    @Override
    public <T> Flux<T> selectList(Class<T> beanClass, String sql, Object... parameters) {
        return fluxFunc(r2dbcConnection -> crud.getCrudSql().selectList(r2dbcConnection, beanClass, sql, parameters));
    }

    @Override
    public Mono<Integer> insert(List<String> pkeys, Map<String, Object> keyValues, String sql, Object... parameters) {
        return monoFunc(r2dbcConnection -> crud.getCrudSql().insert(r2dbcConnection, pkeys, keyValues, sql, parameters));
    }

    @Override
    public Mono<Integer> execute(String sql, Object... parameters) {
        return monoFunc(r2dbcConnection -> crud.getCrudSql().execute(r2dbcConnection, sql, parameters));
    }

    @Override
    public Mono<Integer> execute(String sql) {
        return monoFunc(r2dbcConnection -> crud.getCrudSql().execute(r2dbcConnection, sql));
    }
}
